package llm

import (
	_ "embed"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/go-resty/resty/v2"
)

/* xiaoying is an AI bot designed for translation purposes */

//go:embed xiaoying_system_prompt.md
var xy_system_prompt string

const xy_user_prompt_template = "翻译下面的全部内容：\n%s"

var client = resty.New()

// one token shared between several apps
var qwen_token = &Token{}
var ernie_token = &Token{}

var models = map[string]Model{

	// zhipu GLM-4
	MODEL_GLM_4: &Glm4{},

	// ali QWen [max, plus]
	MODEL_QWEN_MAX:  &Qwen{Name: MODEL_QWEN_MAX, Token: qwen_token},
	MODEL_QWEN_PLUS: &Qwen{Name: MODEL_QWEN_PLUS, Token: qwen_token},

	// baidu ERNIE-BOT & 4.0
	MODEL_ERNIE_BOT:   &ErnieBot{Name: MODEL_ERNIE_BOT, Token: ernie_token},
	MODEL_ERNIE_BOT_4: &ErnieBot{Name: MODEL_ERNIE_BOT_4, Token: ernie_token},

	// xunfei Spark 3.5
	MODEL_SPARK_35: &Spark{},
}

/* return llm original result; if rsp is not nil, unmarshal result into rsp */
func Translate(input string, model_name string, rsp *TranslateRsp) (string, error) {

	model, ok := models[model_name]
	if !ok {
		return "", fmt.Errorf("model %s not found", model_name)
	}

	err := model.RefreshToken()
	if err != nil {
		return "", err
	}

	var start int64
	if rsp != nil {
		start = time.Now().UnixMilli()
	}

	result, err := model.CallApi(xy_system_prompt, fmt.Sprintf(xy_user_prompt_template, input))
	if err != nil {
		return "", err
	}

	if rsp != nil {

		rsp.Elapse = time.Now().UnixMilli() - start
		rsp.Model = model_name

		s := extractJson(result)
		if s == "" {
			return result, errors.New("can't extract json from result, result=" + result)
		}

		err = json.Unmarshal([]byte(s), rsp)
		if err != nil {
			return result, err
		}

	}

	return result, nil
}

func extractJson(c string) string {

	start := strings.IndexByte(c, '{')
	if start == -1 {
		return ""
	}

	end := strings.LastIndexByte(c, '}')
	if end == -1 {
		return ""
	}

	return c[start : end+1]
}
